{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "machine_shape": "hm",
      "gpuType": "A100"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "## Install (Colab Only)"
      ],
      "metadata": {
        "id": "wUs8UtMcEEx5"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "j0YOp_jaDwsM",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "9aca9ac9-7dce-4c91-fbe3-cef05996d4a9"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Cloning into 'PyEPO'...\n",
            "remote: Enumerating objects: 147, done.\u001b[K\n",
            "remote: Counting objects: 100% (147/147), done.\u001b[K\n",
            "remote: Compressing objects: 100% (133/133), done.\u001b[K\n",
            "remote: Total 147 (delta 30), reused 64 (delta 9), pack-reused 0 (from 0)\u001b[K\n",
            "Receiving objects: 100% (147/147), 6.94 MiB | 5.08 MiB/s, done.\n",
            "Resolving deltas: 100% (30/30), done.\n",
            "Processing ./PyEPO/pkg\n",
            "  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.11/dist-packages (from pyepo==1.0.0) (1.26.4)\n",
            "Requirement already satisfied: scipy in /usr/local/lib/python3.11/dist-packages (from pyepo==1.0.0) (1.13.1)\n",
            "Collecting pathos (from pyepo==1.0.0)\n",
            "  Downloading pathos-0.3.3-py3-none-any.whl.metadata (11 kB)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.11/dist-packages (from pyepo==1.0.0) (4.67.1)\n",
            "Collecting configspace (from pyepo==1.0.0)\n",
            "  Downloading configspace-1.2.1.tar.gz (130 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m131.0/131.0 kB\u001b[0m \u001b[31m8.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "  Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "Requirement already satisfied: scikit_learn in /usr/local/lib/python3.11/dist-packages (from pyepo==1.0.0) (1.6.1)\n",
            "Requirement already satisfied: torch>=1.13.1 in /usr/local/lib/python3.11/dist-packages (from pyepo==1.0.0) (2.5.1+cu124)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.11/dist-packages (from torch>=1.13.1->pyepo==1.0.0) (3.17.0)\n",
            "Requirement already satisfied: typing-extensions>=4.8.0 in /usr/local/lib/python3.11/dist-packages (from torch>=1.13.1->pyepo==1.0.0) (4.12.2)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.11/dist-packages (from torch>=1.13.1->pyepo==1.0.0) (3.4.2)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.11/dist-packages (from torch>=1.13.1->pyepo==1.0.0) (3.1.5)\n",
            "Requirement already satisfied: fsspec in /usr/local/lib/python3.11/dist-packages (from torch>=1.13.1->pyepo==1.0.0) (2024.10.0)\n",
            "Collecting nvidia-cuda-nvrtc-cu12==12.4.127 (from torch>=1.13.1->pyepo==1.0.0)\n",
            "  Downloading nvidia_cuda_nvrtc_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl.metadata (1.5 kB)\n",
            "Collecting nvidia-cuda-runtime-cu12==12.4.127 (from torch>=1.13.1->pyepo==1.0.0)\n",
            "  Downloading nvidia_cuda_runtime_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl.metadata (1.5 kB)\n",
            "Collecting nvidia-cuda-cupti-cu12==12.4.127 (from torch>=1.13.1->pyepo==1.0.0)\n",
            "  Downloading nvidia_cuda_cupti_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl.metadata (1.6 kB)\n",
            "Collecting nvidia-cudnn-cu12==9.1.0.70 (from torch>=1.13.1->pyepo==1.0.0)\n",
            "  Downloading nvidia_cudnn_cu12-9.1.0.70-py3-none-manylinux2014_x86_64.whl.metadata (1.6 kB)\n",
            "Collecting nvidia-cublas-cu12==12.4.5.8 (from torch>=1.13.1->pyepo==1.0.0)\n",
            "  Downloading nvidia_cublas_cu12-12.4.5.8-py3-none-manylinux2014_x86_64.whl.metadata (1.5 kB)\n",
            "Collecting nvidia-cufft-cu12==11.2.1.3 (from torch>=1.13.1->pyepo==1.0.0)\n",
            "  Downloading nvidia_cufft_cu12-11.2.1.3-py3-none-manylinux2014_x86_64.whl.metadata (1.5 kB)\n",
            "Collecting nvidia-curand-cu12==10.3.5.147 (from torch>=1.13.1->pyepo==1.0.0)\n",
            "  Downloading nvidia_curand_cu12-10.3.5.147-py3-none-manylinux2014_x86_64.whl.metadata (1.5 kB)\n",
            "Collecting nvidia-cusolver-cu12==11.6.1.9 (from torch>=1.13.1->pyepo==1.0.0)\n",
            "  Downloading nvidia_cusolver_cu12-11.6.1.9-py3-none-manylinux2014_x86_64.whl.metadata (1.6 kB)\n",
            "Collecting nvidia-cusparse-cu12==12.3.1.170 (from torch>=1.13.1->pyepo==1.0.0)\n",
            "  Downloading nvidia_cusparse_cu12-12.3.1.170-py3-none-manylinux2014_x86_64.whl.metadata (1.6 kB)\n",
            "Requirement already satisfied: nvidia-nccl-cu12==2.21.5 in /usr/local/lib/python3.11/dist-packages (from torch>=1.13.1->pyepo==1.0.0) (2.21.5)\n",
            "Requirement already satisfied: nvidia-nvtx-cu12==12.4.127 in /usr/local/lib/python3.11/dist-packages (from torch>=1.13.1->pyepo==1.0.0) (12.4.127)\n",
            "Collecting nvidia-nvjitlink-cu12==12.4.127 (from torch>=1.13.1->pyepo==1.0.0)\n",
            "  Downloading nvidia_nvjitlink_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl.metadata (1.5 kB)\n",
            "Requirement already satisfied: triton==3.1.0 in /usr/local/lib/python3.11/dist-packages (from torch>=1.13.1->pyepo==1.0.0) (3.1.0)\n",
            "Requirement already satisfied: sympy==1.13.1 in /usr/local/lib/python3.11/dist-packages (from torch>=1.13.1->pyepo==1.0.0) (1.13.1)\n",
            "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /usr/local/lib/python3.11/dist-packages (from sympy==1.13.1->torch>=1.13.1->pyepo==1.0.0) (1.3.0)\n",
            "Requirement already satisfied: pyparsing in /usr/local/lib/python3.11/dist-packages (from configspace->pyepo==1.0.0) (3.2.1)\n",
            "Requirement already satisfied: more_itertools in /usr/local/lib/python3.11/dist-packages (from configspace->pyepo==1.0.0) (10.6.0)\n",
            "Collecting ppft>=1.7.6.9 (from pathos->pyepo==1.0.0)\n",
            "  Downloading ppft-1.7.6.9-py3-none-any.whl.metadata (12 kB)\n",
            "Collecting dill>=0.3.9 (from pathos->pyepo==1.0.0)\n",
            "  Downloading dill-0.3.9-py3-none-any.whl.metadata (10 kB)\n",
            "Collecting pox>=0.3.5 (from pathos->pyepo==1.0.0)\n",
            "  Downloading pox-0.3.5-py3-none-any.whl.metadata (8.0 kB)\n",
            "Collecting multiprocess>=0.70.17 (from pathos->pyepo==1.0.0)\n",
            "  Downloading multiprocess-0.70.17-py311-none-any.whl.metadata (7.2 kB)\n",
            "Requirement already satisfied: joblib>=1.2.0 in /usr/local/lib/python3.11/dist-packages (from scikit_learn->pyepo==1.0.0) (1.4.2)\n",
            "Requirement already satisfied: threadpoolctl>=3.1.0 in /usr/local/lib/python3.11/dist-packages (from scikit_learn->pyepo==1.0.0) (3.5.0)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.11/dist-packages (from jinja2->torch>=1.13.1->pyepo==1.0.0) (3.0.2)\n",
            "Downloading nvidia_cublas_cu12-12.4.5.8-py3-none-manylinux2014_x86_64.whl (363.4 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m363.4/363.4 MB\u001b[0m \u001b[31m2.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading nvidia_cuda_cupti_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl (13.8 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m13.8/13.8 MB\u001b[0m \u001b[31m106.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading nvidia_cuda_nvrtc_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl (24.6 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m24.6/24.6 MB\u001b[0m \u001b[31m87.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading nvidia_cuda_runtime_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl (883 kB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m883.7/883.7 kB\u001b[0m \u001b[31m54.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading nvidia_cudnn_cu12-9.1.0.70-py3-none-manylinux2014_x86_64.whl (664.8 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m664.8/664.8 MB\u001b[0m \u001b[31m1.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading nvidia_cufft_cu12-11.2.1.3-py3-none-manylinux2014_x86_64.whl (211.5 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m211.5/211.5 MB\u001b[0m \u001b[31m10.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading nvidia_curand_cu12-10.3.5.147-py3-none-manylinux2014_x86_64.whl (56.3 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m56.3/56.3 MB\u001b[0m \u001b[31m33.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading nvidia_cusolver_cu12-11.6.1.9-py3-none-manylinux2014_x86_64.whl (127.9 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m127.9/127.9 MB\u001b[0m \u001b[31m17.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading nvidia_cusparse_cu12-12.3.1.170-py3-none-manylinux2014_x86_64.whl (207.5 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m207.5/207.5 MB\u001b[0m \u001b[31m3.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading nvidia_nvjitlink_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl (21.1 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m21.1/21.1 MB\u001b[0m \u001b[31m99.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading pathos-0.3.3-py3-none-any.whl (82 kB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m82.1/82.1 kB\u001b[0m \u001b[31m8.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading dill-0.3.9-py3-none-any.whl (119 kB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m119.4/119.4 kB\u001b[0m \u001b[31m12.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading multiprocess-0.70.17-py311-none-any.whl (144 kB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m144.3/144.3 kB\u001b[0m \u001b[31m12.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading pox-0.3.5-py3-none-any.whl (29 kB)\n",
            "Downloading ppft-1.7.6.9-py3-none-any.whl (56 kB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m56.8/56.8 kB\u001b[0m \u001b[31m5.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hBuilding wheels for collected packages: pyepo, configspace\n",
            "  Building wheel for pyepo (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for pyepo: filename=pyepo-1.0.0-py3-none-any.whl size=52112 sha256=2909b8cdecdcd662b63f56c21c533bc35f28fd79586c975aca8ab043000e8d7a\n",
            "  Stored in directory: /tmp/pip-ephem-wheel-cache-od7wcglb/wheels/9e/4b/4c/cd1bd6ec37a2decdd12123b131b5da1062b1492f6d2e55bcda\n",
            "  Building wheel for configspace (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for configspace: filename=configspace-1.2.1-py3-none-any.whl size=115950 sha256=ca7dd65732205358868ebabebc5d40de4a836b7909b259ea8b8b33fdcfe3b32f\n",
            "  Stored in directory: /root/.cache/pip/wheels/11/0f/36/d5027c3eeb038827889830f7efbe6a1bad8956b3eb44ab2f44\n",
            "Successfully built pyepo configspace\n",
            "Installing collected packages: ppft, pox, nvidia-nvjitlink-cu12, nvidia-curand-cu12, nvidia-cufft-cu12, nvidia-cuda-runtime-cu12, nvidia-cuda-nvrtc-cu12, nvidia-cuda-cupti-cu12, nvidia-cublas-cu12, dill, nvidia-cusparse-cu12, nvidia-cudnn-cu12, multiprocess, configspace, pathos, nvidia-cusolver-cu12, pyepo\n",
            "  Attempting uninstall: nvidia-nvjitlink-cu12\n",
            "    Found existing installation: nvidia-nvjitlink-cu12 12.5.82\n",
            "    Uninstalling nvidia-nvjitlink-cu12-12.5.82:\n",
            "      Successfully uninstalled nvidia-nvjitlink-cu12-12.5.82\n",
            "  Attempting uninstall: nvidia-curand-cu12\n",
            "    Found existing installation: nvidia-curand-cu12 10.3.6.82\n",
            "    Uninstalling nvidia-curand-cu12-10.3.6.82:\n",
            "      Successfully uninstalled nvidia-curand-cu12-10.3.6.82\n",
            "  Attempting uninstall: nvidia-cufft-cu12\n",
            "    Found existing installation: nvidia-cufft-cu12 11.2.3.61\n",
            "    Uninstalling nvidia-cufft-cu12-11.2.3.61:\n",
            "      Successfully uninstalled nvidia-cufft-cu12-11.2.3.61\n",
            "  Attempting uninstall: nvidia-cuda-runtime-cu12\n",
            "    Found existing installation: nvidia-cuda-runtime-cu12 12.5.82\n",
            "    Uninstalling nvidia-cuda-runtime-cu12-12.5.82:\n",
            "      Successfully uninstalled nvidia-cuda-runtime-cu12-12.5.82\n",
            "  Attempting uninstall: nvidia-cuda-nvrtc-cu12\n",
            "    Found existing installation: nvidia-cuda-nvrtc-cu12 12.5.82\n",
            "    Uninstalling nvidia-cuda-nvrtc-cu12-12.5.82:\n",
            "      Successfully uninstalled nvidia-cuda-nvrtc-cu12-12.5.82\n",
            "  Attempting uninstall: nvidia-cuda-cupti-cu12\n",
            "    Found existing installation: nvidia-cuda-cupti-cu12 12.5.82\n",
            "    Uninstalling nvidia-cuda-cupti-cu12-12.5.82:\n",
            "      Successfully uninstalled nvidia-cuda-cupti-cu12-12.5.82\n",
            "  Attempting uninstall: nvidia-cublas-cu12\n",
            "    Found existing installation: nvidia-cublas-cu12 12.5.3.2\n",
            "    Uninstalling nvidia-cublas-cu12-12.5.3.2:\n",
            "      Successfully uninstalled nvidia-cublas-cu12-12.5.3.2\n",
            "  Attempting uninstall: nvidia-cusparse-cu12\n",
            "    Found existing installation: nvidia-cusparse-cu12 12.5.1.3\n",
            "    Uninstalling nvidia-cusparse-cu12-12.5.1.3:\n",
            "      Successfully uninstalled nvidia-cusparse-cu12-12.5.1.3\n",
            "  Attempting uninstall: nvidia-cudnn-cu12\n",
            "    Found existing installation: nvidia-cudnn-cu12 9.3.0.75\n",
            "    Uninstalling nvidia-cudnn-cu12-9.3.0.75:\n",
            "      Successfully uninstalled nvidia-cudnn-cu12-9.3.0.75\n",
            "  Attempting uninstall: nvidia-cusolver-cu12\n",
            "    Found existing installation: nvidia-cusolver-cu12 11.6.3.83\n",
            "    Uninstalling nvidia-cusolver-cu12-11.6.3.83:\n",
            "      Successfully uninstalled nvidia-cusolver-cu12-11.6.3.83\n",
            "Successfully installed configspace-1.2.1 dill-0.3.9 multiprocess-0.70.17 nvidia-cublas-cu12-12.4.5.8 nvidia-cuda-cupti-cu12-12.4.127 nvidia-cuda-nvrtc-cu12-12.4.127 nvidia-cuda-runtime-cu12-12.4.127 nvidia-cudnn-cu12-9.1.0.70 nvidia-cufft-cu12-11.2.1.3 nvidia-curand-cu12-10.3.5.147 nvidia-cusolver-cu12-11.6.1.9 nvidia-cusparse-cu12-12.3.1.170 nvidia-nvjitlink-cu12-12.4.127 pathos-0.3.3 pox-0.3.5 ppft-1.7.6.9 pyepo-1.0.0\n"
          ]
        }
      ],
      "source": [
        "# install\n",
        "!pip install pyepo\n",
        "!pip install mpax"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Remove the problematic hook from Google Colab."
      ],
      "metadata": {
        "id": "qUo4slYmERJo"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import sys\n",
        "sys.meta_path = [hook for hook in sys.meta_path if not any(keyword in str(hook) for keyword in [\"google.colab\"])]"
      ],
      "metadata": {
        "id": "_aT8QfY2ERmu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "In this tutorial, we will use MPAX as optimization solver to run PDHG on GPU.\n",
        "\n",
        "[MPAX](https://github.com/MIT-Lu-Lab/MPAX) (Mathematical Programming in JAX) s a highly efficient mathematical programming framework that hardware-accelerated, differentiable, batchable, and distributable, esigned to integrate with modern computational and deep learning workflows.\n",
        "\n",
        "At its core, MPAX leverages PDHG (Primal-Dual Hybrid Gradient), a powerful first-order optimization algorithm particularly efficient for large-scale problems. PDHG is highly compatible with batch processing and GPU acceleration, making it an ideal for deep learning."
      ],
      "metadata": {
        "id": "Y6eDiRx8EUqh"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 1 Large Knapsack on MPAX"
      ],
      "metadata": {
        "id": "9KAHqX7KFFKM"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import pyepo\n",
        "# generate data\n",
        "m = 10000 # number of items\n",
        "k = 50    # resource dimension\n",
        "n = 1000  # number of data\n",
        "caps = [500] * k # capacity\n",
        "p = 5     # feature dimention\n",
        "deg = 4   # polynomial degree\n",
        "e = 0.5   # noise half-width\n",
        "weights, feats, costs = pyepo.data.knapsack.genData(num_data=n+1000, num_features=p, num_items=m,\n",
        "                                                    dim=k, deg=deg, noise_width=e, seed=42)"
      ],
      "metadata": {
        "id": "yPaBnWY9E2Op"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Thus, we can use ``optMpaxModel`` to for linear programming."
      ],
      "metadata": {
        "id": "7pp1bIY9FP9X"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# build optModel\n",
        "from pyepo.model.mpax import knapsackModel\n",
        "optmodel = knapsackModel(weights, caps)"
      ],
      "metadata": {
        "id": "Su_968WdFLs7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 2 Dataset and Data Loader"
      ],
      "metadata": {
        "id": "88hn23gsFt0u"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Similar to other solver, we can use `optDataset`."
      ],
      "metadata": {
        "id": "S7Ib2-oGFvhC"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# split train test data\n",
        "from sklearn.model_selection import train_test_split\n",
        "x_train, x_test, c_train, c_test = train_test_split(feats, costs, test_size=1000, random_state=42)"
      ],
      "metadata": {
        "id": "MdJ9YKJsFo5V"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# get optDataset\n",
        "dataset_train = pyepo.data.dataset.optDataset(optmodel, x_train, c_train)\n",
        "dataset_test = pyepo.data.dataset.optDataset(optmodel, x_test, c_test)"
      ],
      "metadata": {
        "id": "GmdwqvMxGB_P",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "d4382080-c520-4f35-fd49-a165b45135a4"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Optimizing for optDataset...\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 1000/1000 [02:11<00:00,  7.63it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Optimizing for optDataset...\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 1000/1000 [02:07<00:00,  7.86it/s]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# set data loader\n",
        "from torch.utils.data import DataLoader\n",
        "batch_size = 32\n",
        "loader_train = DataLoader(dataset_train, batch_size=batch_size, shuffle=True)\n",
        "loader_test = DataLoader(dataset_test, batch_size=batch_size, shuffle=False)"
      ],
      "metadata": {
        "id": "-6ew31MAGCJN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 3 Linear Regression on PyTorch"
      ],
      "metadata": {
        "id": "fcwo3oeaGTcT"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Here, we build the simplest PyTorch model, linear regression."
      ],
      "metadata": {
        "id": "1Viws-UgGWUv"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from torch import nn\n",
        "# build linear model\n",
        "class LinearRegression(nn.Module):\n",
        "\n",
        "    def __init__(self):\n",
        "        super(LinearRegression, self).__init__()\n",
        "        self.linear = nn.Linear(p, m)\n",
        "\n",
        "    def forward(self, x):\n",
        "        out = self.linear(x)\n",
        "        return out"
      ],
      "metadata": {
        "id": "nQufZn9eGVou"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 3 Training"
      ],
      "metadata": {
        "id": "QiZbtxayGc36"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Define function to train model with different methods."
      ],
      "metadata": {
        "id": "S4mZyosqGxg0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import time\n",
        "\n",
        "# train model\n",
        "def trainModel(reg, loss_func, method_name, num_epochs=10, lr=1e-2):\n",
        "    # set adam optimizer\n",
        "    optimizer = torch.optim.Adam(reg.parameters(), lr=lr)\n",
        "    # train mode\n",
        "    reg.train()\n",
        "    # init log\n",
        "    loss_log = []\n",
        "    loss_log_regret = [pyepo.metric.regret(reg, optmodel, loader_test)]\n",
        "    # init elpased time\n",
        "    elapsed = 0\n",
        "    for epoch in range(num_epochs):\n",
        "        # start timing\n",
        "        tick = time.time()\n",
        "        # load data\n",
        "        for i, data in enumerate(loader_train):\n",
        "            x, c, w, z = data\n",
        "            # cuda\n",
        "            if torch.cuda.is_available():\n",
        "                x, c, w, z = x.cuda(), c.cuda(), w.cuda(), z.cuda()\n",
        "            # forward pass\n",
        "            cp = reg(x)\n",
        "            if method_name == \"spo+\":\n",
        "                loss = loss_func(cp, c, w, z)\n",
        "            if method_name == \"pfy\":\n",
        "                loss = loss_func(cp, w)\n",
        "            if method_name == \"ltr\":\n",
        "                loss = loss_func(cp, c)\n",
        "            # backward pass\n",
        "            optimizer.zero_grad()\n",
        "            loss.backward()\n",
        "            optimizer.step()\n",
        "            # record time\n",
        "            tock = time.time()\n",
        "            elapsed += tock - tick\n",
        "            # log\n",
        "            loss_log.append(loss.item())\n",
        "        regret = pyepo.metric.regret(reg, optmodel, loader_test)\n",
        "        loss_log_regret.append(regret)\n",
        "        print(\"Epoch {:2},  Loss: {:9.4f},  Regret: {:7.4f}%\".format(epoch+1, loss.item(), regret*100))\n",
        "    print(\"Total Elapsed Time: {:.2f} Sec.\".format(elapsed))\n",
        "    return loss_log, loss_log_regret"
      ],
      "metadata": {
        "id": "6Y5oWNUBGkCV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 3.1 Smart Predict-then-Optimize"
      ],
      "metadata": {
        "id": "99wPIYtiHBBm"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# init SPO+ loss\n",
        "spop = pyepo.func.SPOPlus(optmodel)"
      ],
      "metadata": {
        "id": "ujK6Maq4HMCA",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "31c976f7-7732-42c0-98a3-a4f0060e75f4"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Num of cores: 1\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "# init model\n",
        "reg = LinearRegression()\n",
        "# cuda\n",
        "if torch.cuda.is_available():\n",
        "    reg = reg.cuda()\n",
        "# train\n",
        "loss_log, loss_log_regret = trainModel(reg, loss_func=spop, method_name=\"spo+\")"
      ],
      "metadata": {
        "id": "unuoGV6uG9Zw",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "a345dbce-5896-4f2b-ff6d-2b8fa933cee1"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch  1,  Loss: 1051.8181,  Regret: 40.3334%\n",
            "Epoch  2,  Loss: 1556.9424,  Regret: 33.3379%\n",
            "Epoch  3,  Loss: 1159.5161,  Regret: 29.9910%\n",
            "Epoch  4,  Loss:  662.3110,  Regret: 27.9450%\n",
            "Epoch  5,  Loss:  479.3947,  Regret: 26.6849%\n",
            "Epoch  6,  Loss: 1020.9800,  Regret: 25.9281%\n",
            "Epoch  7,  Loss:  645.3926,  Regret: 25.4048%\n",
            "Epoch  8,  Loss:  375.7114,  Regret: 25.1683%\n",
            "Epoch  9,  Loss:  818.3668,  Regret: 24.9954%\n",
            "Epoch 10,  Loss:  800.0767,  Regret: 24.8941%\n",
            "Total Elapsed Time: 1116.01 Sec.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 3.2 Perturbed Fenchel-Young Loss"
      ],
      "metadata": {
        "id": "jClVFJk4HVWY"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# init pfyl loss\n",
        "pfy = pyepo.func.perturbedFenchelYoung(optmodel, n_samples=3, sigma=1.0)"
      ],
      "metadata": {
        "id": "1e0RIBskHXFS",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "03981bbc-534a-4c84-e96f-dd4808f4a1b8"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Num of cores: 1\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "# init model\n",
        "reg = LinearRegression()\n",
        "# cuda\n",
        "if torch.cuda.is_available():\n",
        "    reg = reg.cuda()\n",
        "# train\n",
        "loss_log, loss_log_regret = trainModel(reg, loss_func=pfy, method_name=\"pfy\")"
      ],
      "metadata": {
        "id": "p_MyVKAZHbAJ",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "146a29ad-14e3-44f5-bf3e-9876353f79a5"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch  1,  Loss:  113.6920,  Regret: 43.4046%\n",
            "Epoch  2,  Loss:  110.4038,  Regret: 33.5827%\n",
            "Epoch  3,  Loss:  109.1026,  Regret: 29.4995%\n",
            "Epoch  4,  Loss:  106.5823,  Regret: 27.7579%\n",
            "Epoch  5,  Loss:  106.6420,  Regret: 26.8516%\n",
            "Epoch  6,  Loss:  105.0522,  Regret: 26.3880%\n",
            "Epoch  7,  Loss:  101.1143,  Regret: 26.2023%\n",
            "Epoch  8,  Loss:  100.7843,  Regret: 25.9881%\n",
            "Epoch  9,  Loss:  102.8199,  Regret: 25.7716%\n",
            "Epoch 10,  Loss:  104.2387,  Regret: 25.7466%\n",
            "Total Elapsed Time: 4311.60 Sec.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 3.3 Pointwise Learning To Rank"
      ],
      "metadata": {
        "id": "a6txudR-Hl_Q"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# init ltr loss\n",
        "ptltr = pyepo.func.pointwiseLTR(optmodel, solve_ratio=0.05, dataset=dataset_train)"
      ],
      "metadata": {
        "id": "GGjhrL2DHmd6",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "fb040a5c-314b-426e-d589-215626287ddc"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Num of cores: 1\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "# init model\n",
        "reg = LinearRegression()\n",
        "# cuda\n",
        "if torch.cuda.is_available():\n",
        "    reg = reg.cuda()\n",
        "# train\n",
        "loss_log, loss_log_regret = trainModel(reg, loss_func=ptltr, method_name=\"ltr\")"
      ],
      "metadata": {
        "id": "uNeEZ-YVHmjc",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "28a758a7-be89-4683-bb19-b92c97fa32fb"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch  1,  Loss: 146505.3438,  Regret: 50.8165%\n",
            "Epoch  2,  Loss: 76844.2031,  Regret: 43.2709%\n",
            "Epoch  3,  Loss: 330145.3750,  Regret: 39.5543%\n",
            "Epoch  4,  Loss: 121497.7812,  Regret: 36.8107%\n",
            "Epoch  5,  Loss: 66077.6719,  Regret: 35.3827%\n",
            "Epoch  6,  Loss: 125274.8438,  Regret: 35.1007%\n",
            "Epoch  7,  Loss: 44842.1797,  Regret: 34.6702%\n",
            "Epoch  8,  Loss: 32069.6445,  Regret: 34.1945%\n",
            "Epoch  9,  Loss: 43524.7305,  Regret: 33.7943%\n",
            "Epoch 10,  Loss: 46142.7422,  Regret: 33.2205%\n",
            "Total Elapsed Time: 104.54 Sec.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "htIky8HPHziF"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}